from business.ollama.chat_chat import OllamaChatBase
import asyncio
import json
    
async def test_stream_chat(chat: OllamaChatBase, query: str, session_id: str):
    """测试流式对话"""
    try:
        async for response in chat.stream_chat(query, session_id):
            # 直接打印流式响应的内容
            if response["event"] == "message":
                data = json.loads(response["data"])
            elif response["event"] == "error":
                print(f"\n错误: {response['data']}")
    except Exception as e:
        print(f"\n流式对话发生错误: {e}")

async def main_async():
    # 在测试模式下启用 verbose
    chat = OllamaChatBase(verbose=True)
    
    print("\n请选择测试模式:")
    print("1. 普通对话模式")
    print("2. 流式对话模式")
    mode = input("请输入模式编号 (1 或 2): ").strip()
    
    while True:
        try:
            user_input = input("\n请输入您的问题 (输入 'quit' 退出, 'switch' 切换模式): ")
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("再见！")
                break
            elif user_input.lower() == 'switch':
                mode = '2' if mode == '1' else '1'
                print(f"\n已切换到{'普通' if mode == '1' else '流式'}对话模式")
                continue
                
            if user_input.strip():
                if mode == '1':
                    # 普通对话模式
                    response = chat.chat(user_input)
                    print(response)
                else:
                    # 流式对话模式
                    session_id = "temp_" + str(id('session_id'))
                    await test_stream_chat(chat, user_input, session_id)
                    print()  # 添加一个换行
                
        except KeyboardInterrupt:
            print("\n程序已终止")
            break
        except Exception as e:
            print(f"发生错误: {e}")

def main():
    """入口函数"""
    asyncio.run(main_async())

if __name__ == "__main__":
    main()
